// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/init/init.h>
#include <gtest/gtest.h>
#include <pollux/testing/gtest_utils.h>
#include <pollux/connectors/fuzzer/tests/FuzzerConnectorTestBase.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/plan/plan_builder.h>

namespace kumo::pollux::connector::fuzzer::test {
    class FuzzerConnectorTest : public FuzzerConnectorTestBase {
    };

    using namespace kumo::pollux::plan;

    TEST_F(FuzzerConnectorTest, singleSplit) {
        const size_t numRows = 100;
        auto type = ROW({BIGINT(), DOUBLE(), VARCHAR()});

        auto plan = PlanBuilder()
                .startTableScan()
                .outputType(type)
                .tableHandle(makeFuzzerTableHandle())
                .endTableScan()
                .planNode();

        exec::test::AssertQueryBuilder(plan)
                .split(makeFuzzerSplit(numRows))
                .assertTypeAndNumRows(type, numRows);
    }

    TEST_F(FuzzerConnectorTest, floatingPoints) {
        const size_t numRows = 1000;
        auto type = ROW({REAL(), DOUBLE()});

        auto plan = PlanBuilder()
                .startTableScan()
                .outputType(type)
                .tableHandle(makeFuzzerTableHandle())
                .endTableScan()
                .planNode();

        exec::test::AssertQueryBuilder(plan)
                .split(makeFuzzerSplit(numRows))
                .assertTypeAndNumRows(type, numRows);
    }

    TEST_F(FuzzerConnectorTest, complexTypes) {
        const size_t numRows = 100;
        auto type = ROW({
            ARRAY(BIGINT()),
            ROW({VARCHAR(), MAP(INTEGER(), ARRAY(DOUBLE())), VARBINARY()}),
            REAL(),
        });

        auto plan = PlanBuilder()
                .startTableScan()
                .outputType(type)
                .tableHandle(makeFuzzerTableHandle())
                .endTableScan()
                .planNode();

        exec::test::AssertQueryBuilder(plan)
                .split(makeFuzzerSplit(numRows))
                .assertTypeAndNumRows(type, numRows);
    }

    TEST_F(FuzzerConnectorTest, multipleSplits) {
        const size_t rowsPerSplit = 100;
        const size_t numSplits = 10;
        auto type = ROW({BIGINT(), DOUBLE(), VARCHAR()});

        auto plan = PlanBuilder()
                .startTableScan()
                .outputType(type)
                .tableHandle(makeFuzzerTableHandle())
                .endTableScan()
                .planNode();

        exec::test::AssertQueryBuilder(plan)
                .splits(makeFuzzerSplits(rowsPerSplit, numSplits))
                .assertTypeAndNumRows(type, rowsPerSplit * numSplits);
    }

    TEST_F(FuzzerConnectorTest, randomTypes) {
        const size_t rowsPerSplit = 100;
        const size_t numSplits = 10;

        const size_t iterations = 20;

        for (size_t i = 0; i < iterations; ++i) {
            auto type = VectorFuzzer({}, pool()).randRowType();

            auto plan = PlanBuilder()
                    .startTableScan()
                    .outputType(type)
                    .tableHandle(makeFuzzerTableHandle())
                    .endTableScan()
                    .planNode();
            exec::test::AssertQueryBuilder(plan)
                    .splits(makeFuzzerSplits(rowsPerSplit, numSplits))
                    .assertTypeAndNumRows(type, rowsPerSplit * numSplits);
        }
    }

    TEST_F(FuzzerConnectorTest, reproducible) {
        const size_t numRows = 100;
        auto type = ROW({BIGINT(), ARRAY(INTEGER()), VARCHAR()});

        auto plan1 = PlanBuilder()
                .startTableScan()
                .outputType(type)
                .tableHandle(makeFuzzerTableHandle(/*fuzerSeed=*/1234))
                .endTableScan()
                .planNode();
        auto plan2 = PlanBuilder()
                .startTableScan()
                .outputType(type)
                .tableHandle(makeFuzzerTableHandle(/*fuzerSeed=*/1234))
                .endTableScan()
                .planNode();

        auto results1 = exec::test::AssertQueryBuilder(plan1)
                .split(makeFuzzerSplit(numRows))
                .copyResults(pool());
        auto results2 = exec::test::AssertQueryBuilder(plan2)
                .split(makeFuzzerSplit(numRows))
                .copyResults(pool());

        exec::test::assertEqualResults({results1}, {results2});
    }
} // namespace kumo::pollux::connector::fuzzer::test

int main(int argc, char **argv) {
    testing::InitGoogleTest(&argc, argv);
    melon::Init init(&argc, &argv, false);
    return RUN_ALL_TESTS();
}
